import rclpy
from rclpy.node import Node
from sensor_msgs.msg import LaserScan, PointCloud2, PointField
from nav_msgs.msg import OccupancyGrid
from geometry_msgs.msg import TransformStamped
from std_msgs.msg import Header
import numpy as np
import math
import struct
from scipy.spatial import KDTree
from numpy.linalg import svd
import tf2_ros

class ScanDectRobot(Node):
    def __init__(self):
        super().__init__('scan_2_map_python')
        self.declare_parameter('db_min_value', 5)
        self.declare_parameter('dis_range_value', 10)
        self.subscription = self.create_subscription(
            LaserScan,
            '/scan',
            self.scan_callback,
            50)
        self.subscription
        self.point_cloud_2_publisher = self.create_publisher(PointCloud2, '/pointcloud_out', 10)
        self.publisher = self.create_publisher(LaserScan, '/scan_out', 50)
        self.map_subscription = self.create_subscription(
            OccupancyGrid,
            '/map',
            self.map_callback,
            50)
        self.map_subscription
        self.map_points = None
        self.tf_broadcaster = tf2_ros.TransformBroadcaster(self)

    def scan_callback(self, msg):
        if self.map_points is None or not self.map_points.any():
            return
        coordinate_points=[]
        for i, r in enumerate(msg.ranges):
            # 如果超出雷达范围，设置为'nan'
            if r < msg.range_min or r > msg.range_max or math.isnan(r):
                continue
            angle = msg.angle_min + i * msg.angle_increment
            # 将极坐标转换为直角坐标系下的位置
            x = r * math.cos(angle)
            y = r * math.sin(angle)
            coordinate_points.append((x, y))

        # 发布点云进行debug
        # self.pub_point_cloud(coordinate_points)
            
        # 将数组转为矩阵
        scan_points=self.arr_to_np(coordinate_points)
        best_x, best_y, best_theta, best_error = self.brute_force_match(self.map_points, scan_points)
        print('best_x: ', best_x)
        print('best_y: ', best_y)
        print('best_theta: ', best_theta)
        print('best_error: ', best_error)
        # 可视化结果（将2号机器人的激光雷达数据旋转后与1号机器人的激光雷达数据对齐）
        rotation_matrix = np.array([
            [np.cos(best_theta), -np.sin(best_theta)],
            [np.sin(best_theta), np.cos(best_theta)]
        ])
        scan_points_rough = np.matmul(rotation_matrix, scan_points.T).T
        scan_points_rough[:, 0] += best_x
        scan_points_rough[:, 1] += best_y
        # 记录累积的R和t
        R_cum = rotation_matrix
        t_cum = np.array([best_x, best_y]).T
        # 精确位姿匹配
        scan_points_precise = scan_points_rough
        # 使用icp算法进行精确位姿匹配
        for i in range(100):
            R, t = self.icp(scan_points_precise, self.map_points)
            scan_points_precise = np.matmul(R, scan_points_precise.T).T + t
            # 计算误差
            kdtree = KDTree(self.map_points)
            distances, indices = kdtree.query(scan_points_precise)
            error = np.sum(distances) / distances.shape[0]
            # 更新累积的R和t
            R_cum = np.matmul(R, R_cum)
            t_cum = np.matmul(R, t_cum) + t
            print('error: ', error)
            if error < 0.05:
                print('converged!')
                break
        # R转换为角度
        theta = np.arctan2(R_cum[1, 0], R_cum[0, 0])
        print('theta: ', theta)
        print('t_cum: ', t_cum)
        transform_stamped = TransformStamped()
        transform_stamped.header.stamp = self.get_clock().now().to_msg()
        transform_stamped.header.frame_id = 'map'
        transform_stamped.child_frame_id = 'odom'

        # 设置 map 到 odom 的变换关系，假设没有变换
        transform_stamped.transform.translation.x = t_cum[0]
        transform_stamped.transform.translation.y = t_cum[1]
        transform_stamped.transform.translation.z = 0.0

        quaternion = self.euler_to_quaternion(0,0,theta)
        transform_stamped.transform.rotation.x = quaternion[0]
        transform_stamped.transform.rotation.y = quaternion[1]
        transform_stamped.transform.rotation.z = quaternion[2]
        transform_stamped.transform.rotation.w = quaternion[3]
        self.tf_broadcaster.sendTransform(transform_stamped)
        


    def map_callback(self, msg):
        print('map is ok')
        map_data = np.array(msg.data)
        map_info = msg.info
        # 将地图上的占用点存到数组中
        non_zero_points = []
        for row in range(map_info.height):
            for col in range(map_info.width):
                if map_data[row * map_info.width + col] != 0:
                    x = map_info.origin.position.x + col * map_info.resolution
                    y = map_info.origin.position.y + row * map_info.resolution
                    non_zero_points.append((x, y))
        
        # 发布点云
        # self.pub_point_cloud(non_zero_points)

        # 将数组转为矩阵
        self.map_points=self.arr_to_np(non_zero_points)
        
    # 将数组转为矩阵
    def arr_to_np(self,points):
        np_points = np.zeros((len(points), 2), dtype=np.float32)
        for i in range(len(points)):
            np_points[i][0] = points[i][0]
            np_points[i][1] = points[i][1]
        return np_points

    # 输入点集，发布点云
    def pub_point_cloud(self,points):
        header = Header()
        header.frame_id = "laser"
        point_cloud_data = []
        for point in points:
            point_cloud_data.extend([point[0], point[1], 0.0]) # x, y, z=0
        
        pc2 = PointCloud2()
        pc2.header = header
        pc2.height = 1
        pc2.width = len(points)
        pc2.fields = [
            PointField(name='x', offset=0, datatype=PointField.FLOAT32, count=1),
            PointField(name='y', offset=4, datatype=PointField.FLOAT32, count=1),
            PointField(name='z', offset=8, datatype=PointField.FLOAT32, count=1)
        ]
        pc2.point_step = 12
        pc2.row_step = pc2.point_step * len(points)
        pc2.is_bigendian = False
        pc2.is_dense = True
        pc2.data = struct.pack('<{}f'.format(len(point_cloud_data)), *point_cloud_data)
        
        # Publish the PointCloud2 message
        self.point_cloud_2_publisher.publish(pc2)

    # 暴搜位姿匹配
    def brute_force_match(self,lidar_coords_1, lidar_coords_2):
        # 两个位姿的x,y,theta的搜索范围
        x_range = np.arange(-10, 10, 2)
        y_range = np.arange(-10, 10, 2)
        theta_range = np.arange(-np.pi, np.pi, np.pi / 5)
        # 最佳匹配位姿
        best_x = 0
        best_y = 0
        best_theta = 0
        # 最佳匹配位姿的误差
        best_error = 100000000
        # 搜索
        for x in x_range:
            for y in y_range:
                for theta in theta_range:
                    # 旋转矩阵
                    rotation_matrix = np.array([
                        [np.cos(theta), -np.sin(theta)],
                        [np.sin(theta), np.cos(theta)]
                    ])
                    # 旋转后的激光雷达数据
                    lidar_coords_2_transformed = np.matmul(rotation_matrix, lidar_coords_2.T).T
                    # 平移后的激光雷达数据
                    lidar_coords_2_transformed[:, 0] += x
                    lidar_coords_2_transformed[:, 1] += y
                    # 误差（最近邻距离的和，使用KDTree加速）
                    error = 0
                    kdtree = KDTree(lidar_coords_1)
                    distances, indices = kdtree.query(lidar_coords_2_transformed)
                    # 计算误差
                    error = np.sum(distances) / distances.shape[0]                
                    # 更新最佳匹配位姿
                    if error < best_error:
                        best_error = error
                        best_x = x
                        best_y = y
                        best_theta = theta
        return best_x, best_y, best_theta, best_error

    def icp(self,source, target):
        # 使用kdtree搜索最近邻作为匹配点
        kdtree = KDTree(target)
        distances, indices = kdtree.query(source)
        target = target[indices]
        # 初始化R和t
        R = np.eye(3)
        t = np.zeros((1, 3))
        # 去中心化
        source_mean = np.mean(source, axis=0)
        target_mean = np.mean(target, axis=0)
        source = source - source_mean
        target = target - target_mean
        # 计算协方差矩阵
        H = np.matmul(source.T, target)
        # 使用SVD分解计算旋转矩阵
        U, S, V = np.linalg.svd(H)
        R = np.matmul(V.T, U.T)
        # 计算平移矩阵
        t = target_mean.T - np.matmul(R, source_mean.T)
        return R, t

    # 欧拉角转四元数
    def euler_to_quaternion(self,roll, pitch, yaw):
        cy = np.cos(yaw * 0.5)
        sy = np.sin(yaw * 0.5)
        cp = np.cos(pitch * 0.5)
        sp = np.sin(pitch * 0.5)
        cr = np.cos(roll * 0.5)
        sr = np.sin(roll * 0.5)

        qw = cy * cp * cr + sy * sp * sr
        qx = cy * cp * sr - sy * sp * cr
        qy = sy * cp * sr + cy * sp * cr
        qz = sy * cp * cr - cy * sp * sr

        return np.array([qx, qy, qz, qw])
def main(args=None):
    rclpy.init(args=args)
    node = ScanDectRobot()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
